from sklearn.model_selection import cross_val_score
from sklearn.tree import DecisionTreeRegressor
import numpy as np
import matplotlib.pyplot as plt

rng = np.random.RandomState(1) #生成随机数种子
X = np.sort(5 * rng.rand(80,1), axis=0)
y = np.sin(X).ravel() #降为1维
y[::5] += 3 * (0.5 - rng.rand(16))  #添加噪声


regr_1 = DecisionTreeRegressor(max_depth=5)
regr_2 = DecisionTreeRegressor(max_depth=2)

regr_1.fit(X,y)
regr_2.fit(X,y)

X_test = np.arange(0.0,5.0,0.01)[:,np.newaxis] #测试集数据，增维,输入模型的数据不能是一维的
y_1 = regr_1.predict(X_test)
y_2 = regr_2.predict(X_test)
plt.figure()
plt.scatter(X,y,s=20,edgecolors="black",c="darkorange",label="data")
plt.plot(X_test,y_1,color="cornflowerblue",label="max_depth=5",linewidth=2)
plt.plot(X_test,y_2,color="yellowgreen",label="max_depth=2",linewidth=2)
plt.xlabel("data")
plt.ylabel("target")
plt.title("Decision Tree Regression")
plt.legend()
plt.show()